import torch
import torchvision
from torch import nn
from d2l import torch as d2l

train_augs = torchvision.transforms.Compose([
    torchvision.transforms.RandomVerticalFlip(),
    torchvision.transforms.ToTensor()
])
test_augs = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])

def load_cifar10(is_train,augs,batch_size):
    dataset = torchvision.datasets.CIFAR10(
        train=True,root='DL/CNN/dataset',transform=augs,download=True)
    d2l.show_images([dataset[i][0] for i in range(32)],4,8,scale=0.8)
    d2l.plt.show()
    data_loader = torch.utils.data.DataLoader(
        dataset,batch_size=batch_size,shuffle=is_train,num_workers=4)
    return data_loader

def train_batch(net,X,y,loss,trainer,devices):
    if isinstance(X,list):
        X = [x.to(devices[0]) for x in X]
    else:
        X = X.to(devices[0])
    y = y.to(devices[0])
    net.train()
    trainer.zero_grad()
    y_hat = net(X)
    l = loss(y_hat,y)
    l.sum().backward()
    trainer.step()
    train_loss_sum = l.sum()
    train_acc_sum = d2l.accuracy(y_hat,y)
    return train_loss_sum,train_acc_sum




